#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
人脸识别核心模块 - 主要的人脸识别系统类
"""

import numpy as np
import cv2
import face_recognition
import time
import os
from collections import deque

from config import *
from serial_utils import SerialManager
from utils import *


class FaceRecognitionSystem:
    """OV7670人脸识别系统 - WSL2版本"""
    
    def __init__(self, port=DEFAULT_PORT, baudrate=DEFAULT_BAUDRATE, debug_mode=False):
        """
        初始化OV7670人脸识别系统 - WSL2版本
        
        Args:
            port: 串口设备路径
            baudrate: 波特率
            debug_mode: 调试模式，True时从test文件夹读取图片进行测试
        """
        self.port = port
        self.baudrate = baudrate
        self.debug_mode = debug_mode
        self.serial_manager = SerialManager(port, baudrate)
        self.running = False
        self.frame_count = 0
        
        # 图像参数 - 分离原始数据尺寸和显示尺寸
        self.original_width = DEFAULT_WIDTH    # 原始数据宽度（用于接收数据）
        self.original_height = DEFAULT_HEIGHT  # 原始数据高度（用于接收数据）
        self.width = DEFAULT_WIDTH             # 当前显示宽度
        self.height = DEFAULT_HEIGHT           # 当前显示高度
        self.update_image_size(show_message=False)
        
        # 保存最新接收到的图片（已旋转和处理过的）
        self.latest_frame = None
        self.frame_updated = False  # 标记是否有新图片
        
        # 人脸识别相关
        self.known_face_encodings = []
        self.known_face_names = []
        self.face_locations = []
        self.face_encodings = []
        self.face_names = []
        self.process_this_frame = True
        
        # 性能统计
        self.recognition_times = deque(maxlen=10)
        
        # 人脸报警相关
        self.last_alarm_time = 0
        self.alarm_interval = ALARM_INTERVAL
        self.alarm_enabled = True
        
        # WSL2特定设置
        self.is_wsl = check_wsl_environment()
        setup_display_environment()
        
        # Debug模式设置
        if self.debug_mode:
            print("🐛 Debug模式已启用 - 将从test文件夹读取图片进行测试")
            self.current_test_image_index = 0
    
    def update_image_size(self, show_message=True):
        """更新图像大小参数 - 基于原始尺寸"""
        self.image_size = self.original_width * self.original_height * RGB565_BYTES_PER_PIXEL
        if show_message:
            print(f"📐 数据接收尺寸: {self.original_width}x{self.original_height}, 数据大小: {self.image_size} 字节")
    
    def set_image_dimensions(self, width, height):
        """
        设置原始图像尺寸（用于数据接收）
        
        Args:
            width: 原始图像宽度
            height: 原始图像高度
        """
        if width != self.original_width or height != self.original_height:
            self.original_width = width
            self.original_height = height
            self.update_image_size()
    
    def load_known_faces(self, faces_dir=KNOWN_FACES_DIR):
        """
        从指定目录加载已知人脸图像
        
        Args:
            faces_dir: 存放已知人脸图像的目录
        """
        if not ensure_directory_exists(faces_dir):
            print(f"请将已知人脸图像放入 {faces_dir} 目录中")
            print("图像文件名将作为人名显示")
            return
        
        face_files = get_images_from_directory(faces_dir)
        
        if not face_files:
            print(f"在 {faces_dir} 目录中未找到人脸图像文件")
            print(f"支持的格式: {', '.join(SUPPORTED_IMAGE_FORMATS)}")
            return
        
        print(f"正在加载已知人脸图像...")
        
        for face_file in face_files:
            try:
                # 提取文件名作为人名（去掉扩展名）
                name = os.path.splitext(face_file)[0]
                
                # 加载图像
                image_path = os.path.join(faces_dir, face_file)
                image = face_recognition.load_image_file(image_path)
                
                # 获取人脸编码
                face_encodings = face_recognition.face_encodings(image)
                
                if face_encodings:
                    # 使用第一个检测到的人脸
                    face_encoding = face_encodings[0]
                    self.known_face_encodings.append(face_encoding)
                    self.known_face_names.append(name)
                    print(f"  ✓ 已加载: {name}")
                else:
                    print(f"  ✗ 未检测到人脸: {face_file}")
                    
            except Exception as e:
                print(f"  ✗ 加载失败 {face_file}: {e}")
        
        print(f"共加载 {len(self.known_face_encodings)} 个已知人脸")
    
    def get_test_images(self):
        """获取test文件夹中的测试图像列表"""
        if not ensure_directory_exists(TEST_IMAGES_DIR):
            print(f"请将测试图像放入 {TEST_IMAGES_DIR} 目录中")
            return []
        
        test_files = get_images_from_directory(TEST_IMAGES_DIR)
        
        if not test_files:
            print(f"在 {TEST_IMAGES_DIR} 目录中未找到测试图像文件")
            print(f"支持的格式: {', '.join(SUPPORTED_IMAGE_FORMATS)}")
            return []
        
        return test_files
    
    def load_next_test_image(self):
        """加载下一张测试图像"""
        test_files = self.get_test_images()
        
        if not test_files:
            return None
        
        # 循环读取图像
        if self.current_test_image_index >= len(test_files):
            self.current_test_image_index = 0
        
        image_file = test_files[self.current_test_image_index]
        image_path = os.path.join(TEST_IMAGES_DIR, image_file)
        
        try:
            # 使用OpenCV加载图像
            image = cv2.imread(image_path)
            if image is None:
                print(f"无法读取图像: {image_path}")
                self.current_test_image_index += 1
                return None
            
            # 转换为RGB格式（face_recognition库使用RGB）
            rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            
            original_height, original_width = rgb_image.shape[:2]
            
            if self.debug_mode:
                # Debug模式：保持原始图像尺寸，不进行resize
                # 更新系统的原始图像尺寸参数以匹配当前图像
                self.set_image_dimensions(original_width, original_height)
                # 设置显示尺寸为图像尺寸
                self.width = original_width
                self.height = original_height
                print(f"🖼️  加载测试图像: {image_file} ({self.current_test_image_index + 1}/{len(test_files)}) - 尺寸: {original_width}x{original_height} (保持原始尺寸)")
            else:
                # 非Debug模式：调整到相机分辨率
                rgb_image = cv2.resize(rgb_image, (self.original_width, self.original_height))
                # 显示尺寸与原始尺寸相同（未旋转时）
                self.width = self.original_width 
                self.height = self.original_height
                print(f"🖼️  加载测试图像: {image_file} ({self.current_test_image_index + 1}/{len(test_files)}) - 原始: {original_width}x{original_height} → 调整为: {self.original_width}x{self.original_height}")
            
            self.current_test_image_index += 1
            
            # 保存最新加载的测试图片
            self.latest_frame = rgb_image
            self.frame_updated = True
            
            return rgb_image
            
        except Exception as e:
            print(f"加载测试图像失败 {image_file}: {e}")
            self.current_test_image_index += 1
            return None
    
    def receive_frame(self):
        """接收一帧图像数据 - 完整协议处理"""
        if not self.serial_manager.is_connected():
            return None
        
        try:
            # 1. 等待帧开始标识
            while self.running:
                line = self.serial_manager.read_line_timeout()
                if line and line.startswith('FRAME_START:'):
                    frame_num = int(line.split(':')[1])
                    print(f"开始接收第 {frame_num} 帧")
                    break
            else:
                return None
            
            # 4. 等待数据开始标记
            data_start = self.serial_manager.read_line_timeout()
            if not data_start or data_start != 'DATA_START':
                print(f"数据开始标记错误: {data_start}")
                return None
            
            # 5. 读取二进制图像数据 (提高效率)
            print(f"开始读取 {self.image_size} 字节的图像数据...")
            image_data = b''
            bytes_to_read = self.image_size
            
            start_time = time.time()
            chunk_size = 4096  # 增大读取块大小提高效率
            
            while len(image_data) < self.image_size and time.time() - start_time < 8.0:
                remaining = self.image_size - len(image_data)
                read_size = min(chunk_size, remaining)
                
                if self.serial_manager.serial_conn.in_waiting > 0:
                    chunk = self.serial_manager.read_bytes(read_size)
                    if chunk:
                        image_data += chunk
                        
                        # 每接收一定数据量显示进度
                        if len(image_data) % 30000 == 0:
                            progress = (len(image_data) / self.image_size) * 100
                            print(f"接收进度: {progress:.1f}%")
            
            if len(image_data) != self.image_size:
                print(f"图像数据大小错误: 期望 {self.image_size}, 实际 {len(image_data)}")
                return None
            
            print("图像数据接收完成")
            
            # 6. 等待数据结束标记
            while self.running:
                line = self.serial_manager.read_line_timeout(1.0)  # 较短超时
                if line == 'DATA_END':
                    print("找到DATA_END标记")
                    break
                elif line and 'DATA_END' in line:
                    print(f"找到包含DATA_END的行: {line}")
                    break
            
            # 7. 等待帧结束标记
            frame_end = self.serial_manager.read_line_timeout(1.0)
            if frame_end and frame_end.startswith('FRAME_END:'):
                end_frame_num = int(frame_end.split(':')[1])
                if end_frame_num == frame_num:
                    print(f"第 {frame_num} 帧接收完成 ✓")
                else:
                    print(f"帧号不匹配: 开始{frame_num} vs 结束{end_frame_num}")
            else:
                print(f"帧结束标记异常: {frame_end}")
            
            # 转换图像格式 (RGB565 -> RGB888) - 使用原始尺寸
            rgb888_image = rgb565_to_rgb888(image_data, self.original_width, self.original_height)
            
            if rgb888_image is not None:
                # 根据配置进行图像旋转修正
                if IMAGE_ROTATION_ENABLED:
                    if IMAGE_ROTATION_CLOCKWISE:
                        # 顺时针旋转90度修正图像方向
                        rgb888_image = np.rot90(rgb888_image, k=-1)  # k=-1表示顺时针90度
                        print(f"🔄 图像已顺时针旋转{IMAGE_ROTATION_DEGREES}度")
                    else:
                        # 逆时针旋转90度
                        rgb888_image = np.rot90(rgb888_image, k=1)   # k=1表示逆时针90度
                        print(f"🔄 图像已逆时针旋转{IMAGE_ROTATION_DEGREES}度")
                    
                    # 确保旋转后的数组是连续的
                    if not rgb888_image.flags['C_CONTIGUOUS']:
                        rgb888_image = np.ascontiguousarray(rgb888_image)
                    
                    # 更新显示尺寸（旋转后宽高互换）
                    rotated_height, rotated_width = rgb888_image.shape[:2]
                    self.width = rotated_width
                    self.height = rotated_height
                    print(f"📐 显示尺寸已更新: {self.original_width}x{self.original_height} → {self.width}x{self.height}")
                else:
                    # 没有旋转时，显示尺寸与原始尺寸相同
                    self.width = self.original_width
                    self.height = self.original_height
                
                # 验证最终图像数据质量
                if rgb888_image.size > 0:
                    # 检查图像数据范围
                    min_val, max_val = rgb888_image.min(), rgb888_image.max()
                    if max_val - min_val < 10:
                        print(f"⚠️ 图像对比度较低: 范围[{min_val}, {max_val}]")
                    
                    # 确保数据类型正确
                    if rgb888_image.dtype != np.uint8:
                        rgb888_image = rgb888_image.astype(np.uint8)
                        print("🔧 已修正图像数据类型为uint8")
                else:
                    print("❌ 图像数据为空")
                    return None
            
            # 更新帧计数
            self.frame_count += 1
            
            # 保存最新接收到的图片
            self.latest_frame = rgb888_image
            self.frame_updated = True
            print(f"💾 已保存最新图片到缓存")
            
            # 只在成功时显示简短信息
            elapsed_time = time.time() - start_time
            speed = len(image_data) / elapsed_time / 1024 if elapsed_time > 0 else 0
            print(f"✅ 帧{frame_num} ({speed:.0f}KB/s)")
            
            return rgb888_image
            
        except Exception as e:
            print(f"❌ 接收失败: {e}")
            return None
    
    def recognize_faces(self, frame=None):
        """
        在图像中识别人脸
        
        Args:
            frame: RGB图像数组（可选，如果为None则使用最新保存的图片）
            
        Returns:
            处理后的BGR图像（用于OpenCV显示）
        """
        recognition_start = time.time()
        
        # 优先使用最新保存的图片，如果没有则使用传入的frame
        current_frame = self.latest_frame if self.latest_frame is not None else frame
        
        if current_frame is None:
            print("⚠️ 没有可用的图像数据")
            # 创建一个黑色的错误图像
            bgr_frame = np.zeros((240, 320, 3), dtype=np.uint8)
            cv2.putText(bgr_frame, "No Image Available", (30, 120), 
                       cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
            return bgr_frame
        
        # 只在每隔一帧处理人脸识别以提高性能
        if self.process_this_frame:
            # 确保图像数据格式正确
            rgb_frame = current_frame
            
            # 验证图像数据有效性
            if rgb_frame is None or rgb_frame.size == 0:
                print("⚠️ 无效的图像数据")
                self.face_locations = []
                self.face_encodings = []
                self.face_names = []
            else:
                # 确保图像数据连续性和正确的数据类型
                if not rgb_frame.flags['C_CONTIGUOUS']:
                    rgb_frame = np.ascontiguousarray(rgb_frame)
                
                # 确保数据类型为uint8
                if rgb_frame.dtype != np.uint8:
                    rgb_frame = rgb_frame.astype(np.uint8)
                
                # 确保图像尺寸合理（最小尺寸检查）
                if rgb_frame.shape[0] < 50 or rgb_frame.shape[1] < 50:
                    print(f"⚠️ 图像尺寸过小: {rgb_frame.shape}")
                    self.face_locations = []
                    self.face_encodings = []
                    self.face_names = []
                else:
                    try:
                        # 检测人脸位置和编码
                        self.face_locations = face_recognition.face_locations(rgb_frame, model=FACE_LOCATIONS_MODEL)
                        self.face_encodings = face_recognition.face_encodings(rgb_frame, self.face_locations)
                        
                        self.face_names = []
                        for face_encoding in self.face_encodings:
                            # 与已知人脸进行比较
                            matches = face_recognition.compare_faces(self.known_face_encodings, face_encoding, tolerance=FACE_RECOGNITION_TOLERANCE)
                            name = "Unknown"
                            
                            # 如果有已知人脸数据，找到最匹配的
                            if self.known_face_encodings:
                                face_distances = face_recognition.face_distance(self.known_face_encodings, face_encoding)
                                best_match_index = np.argmin(face_distances)
                                if matches[best_match_index]:
                                    name = self.known_face_names[best_match_index]
                            
                            self.face_names.append(name)
                            
                    except Exception as e:
                        print(f"❌ 人脸识别处理错误: {e}")
                        self.face_locations = []
                        self.face_encodings = []
                        self.face_names = []
        
        self.process_this_frame = not self.process_this_frame
        
        # 转换为BGR格式用于OpenCV显示 - 使用当前帧
        if current_frame is not None and current_frame.size > 0:
            # 确保原始frame数据也是连续的
            if not current_frame.flags['C_CONTIGUOUS']:
                current_frame = np.ascontiguousarray(current_frame)
            
            bgr_frame = cv2.cvtColor(current_frame, cv2.COLOR_RGB2BGR)
        else:
            # 创建一个黑色的错误图像
            bgr_frame = np.zeros((240, 320, 3), dtype=np.uint8)
            cv2.putText(bgr_frame, "Image Error", (50, 120), 
                       cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
        
        # 在图像上绘制人脸框和名称（使用原始坐标，无需缩放）
        for (top, right, bottom, left), name in zip(self.face_locations, self.face_names):
            # 绘制人脸框（增加线条粗细）
            cv2.rectangle(bgr_frame, (left, top), (right, bottom), COLOR_RED, 4)
            
            # 计算标签尺寸（增大字体）
            font = cv2.FONT_HERSHEY_DUPLEX
            font_scale = FONT_SCALE_LARGE
            thickness = FONT_THICKNESS_THICK
            (text_width, text_height), baseline = cv2.getTextSize(name, font, font_scale, thickness)
            
            # 标签背景位置（在人脸框上方）
            label_top = top - text_height - 15
            label_bottom = top - 5
            
            # 确保标签不会超出图像顶部
            if label_top < 0:
                label_top = bottom + 10
                label_bottom = bottom + text_height + 20
            
            # 绘制姓名标签背景（增加背景大小）
            cv2.rectangle(bgr_frame, (left, label_top), (left + text_width + 20, label_bottom), COLOR_RED, cv2.FILLED)
            
            # 绘制姓名文字
            cv2.putText(bgr_frame, name, (left + 10, label_bottom - 10), font, font_scale, COLOR_WHITE, thickness)
        
        # 记录识别时间
        recognition_time = time.time() - recognition_start
        self.recognition_times.append(recognition_time)
        
        # 显示性能信息（调整为更小字体和左上角位置）
        if len(self.recognition_times) > 0:
            avg_time = sum(self.recognition_times) / len(self.recognition_times)
            fps = 1.0 / avg_time if avg_time > 0 else 0
            
            # 使用更小的字体大小
            info_font_scale = 0.6  # 缩小字体
            info_thickness = 1     # 减小线条粗细
            
            # 在图像上显示检测到的人脸数（调整到左上角）
            cv2.putText(bgr_frame, f"Faces: {len(self.face_locations)}", (10, 20), 
                       cv2.FONT_HERSHEY_SIMPLEX, info_font_scale, COLOR_GREEN, info_thickness)
            
            # 显示当前图像尺寸
            cv2.putText(bgr_frame, f"Size: {self.width}x{self.height}", (10, 40), 
                       cv2.FONT_HERSHEY_SIMPLEX, info_font_scale, COLOR_BLUE, info_thickness)
            
            # 显示WSL2标识
            if self.is_wsl:
                cv2.putText(bgr_frame, "WSL2", (10, 60), 
                           cv2.FONT_HERSHEY_SIMPLEX, info_font_scale, COLOR_YELLOW, info_thickness)
                           
            # 显示报警状态
            if self.alarm_enabled:
                # 报警功能开启时，根据是否检测到人脸显示动态状态
                face_detected = len(self.face_locations) > 0
                alarm_status = "ON" if face_detected else "OFF"
                alarm_color = COLOR_GREEN if face_detected else COLOR_RED
            else:
                # 报警功能关闭时显示DISABLED
                alarm_status = "DISABLED"
                alarm_color = COLOR_YELLOW
            
            alarm_y_pos = 80 if self.is_wsl else 60
            cv2.putText(bgr_frame, f"Alarm: {alarm_status}", (10, alarm_y_pos), 
                       cv2.FONT_HERSHEY_SIMPLEX, info_font_scale, alarm_color, info_thickness)
        
        # 人脸检测报警逻辑
        if len(self.face_locations) > 0 and not self.debug_mode:
            # 统计已知和未知人脸
            known_count = sum(1 for name in self.face_names if name != "Unknown")
            unknown_count = sum(1 for name in self.face_names if name == "Unknown")
            total_faces = len(self.face_locations)
            
            # 发送报警数据到STM32
            self.send_face_alarm(total_faces, known_count, unknown_count)
        
        return bgr_frame
    
    def send_face_alarm(self, face_count, known_faces, unknown_faces):
        """
        发送人脸检测报警标志位到STM32下位机
        
        Args:
            face_count: 检测到的人脸总数
            known_faces: 已知人脸数量
            unknown_faces: 未知人脸数量
        """
        if not self.alarm_enabled or not self.serial_manager.is_connected():
            return
        
        current_time = time.time()
        
        # 检查报警间隔，避免过于频繁
        if current_time - self.last_alarm_time < self.alarm_interval:
            return
        
        try:
            # 发送简单的报警标志位：1表示检测到人脸
            success = self.serial_manager.send_data(ALARM_SIGNAL)
            
            if success:
                # 更新最后报警时间
                self.last_alarm_time = current_time
                
                # 打印报警信息
                print(f"\n🚨 报警信号已发送: 检测到 {face_count} 张人脸")
            
        except Exception as e:
            print(f"\n报警发送失败: {e}")
    
    def toggle_alarm(self):
        """切换报警功能开关"""
        self.alarm_enabled = not self.alarm_enabled
        status = "启用" if self.alarm_enabled else "禁用"
        print(f"\n报警功能已{status}")
        return self.alarm_enabled
    
    def get_latest_frame(self):
        """获取最新的图片"""
        return self.latest_frame
    
    def has_new_frame(self):
        """检查是否有新图片"""
        return self.frame_updated
    
    def mark_frame_processed(self):
        """标记图片已处理"""
        self.frame_updated = False 